ANDROID: KVM: arm64: Fix memory ordering for pKVM module callbacks

Registration of module callbacks for the pKVM hypervisor is lockless
thanks to the use of a cmpxchg.

Problem, a CPU can speculatively execute an indirect branch and
speculatively read variables used in that branch. We then need to order
the memory access between variables potentially set in the driver init
(before the callback registration happen) and the call to that
registered callback.

e.g. in the case of the serial.

 CPU0:                                   CPU1:

   driver_init():                        hyp_serial_enabled()
     base_addr = 0xdeadbeef;               enabled = __hyp_putc
     barrier();                            barrier();
     ops->register_serial_driver(putc);    if (enabled)
                                                __hyp_putc(); /* read base_addr */

This is the same for the SMC and PSCI handler callbacks. The abort and
fault callbacks are not impacted: the driver init can only happen before
the kernel is deprivileged i.e. before the host stage-2 is in place and
then before any of those callbacks can be triggered.

Instead of a full barrier, we can use the acquire/release semantics:
relaxing cmpxchg to cmpxchg_release in the registration path and use a
load_acquire in hyp_serial_enabled().

Bug: 292470326
Change-Id: I4b5fe3713fe40cc5ab42ea0e9cdf54e8315dfb44
Signed-off-by: Vincent Donnefort <vdonnefort@google.com>
diff --git a/arch/arm64/kvm/hyp/nvhe/hyp-main.c b/arch/arm64/kvm/hyp/nvhe/hyp-main.c
index 16abb1e..48e20ad 100644
--- a/arch/arm64/kvm/hyp/nvhe/hyp-main.c
+++ b/arch/arm64/kvm/hyp/nvhe/hyp-main.c
@@ -39,7 +39,12 @@
 
 int __pkvm_register_host_smc_handler(bool (*cb)(struct kvm_cpu_context *))
 {
-	return cmpxchg(&default_host_smc_handler, NULL, cb) ? -EBUSY : 0;
+	/*
+	 * Paired with smp_load_acquire(&default_host_smc_handler) in
+	 * handle_host_smc(). Ensure memory stores happening during a pKVM module
+	 * init are observed before executing the callback.
+	 */
+	return cmpxchg_release(&default_host_smc_handler, NULL, cb) ? -EBUSY : 0;
 }
 
 int __pkvm_register_default_trap_handler(bool (*cb)(struct kvm_cpu_context *))
@@ -1376,7 +1381,7 @@
 	handled = kvm_host_psci_handler(host_ctxt);
 	if (!handled)
 		handled = kvm_host_ffa_handler(host_ctxt);
-	if (!handled && READ_ONCE(default_host_smc_handler))
+	if (!handled && smp_load_acquire(&default_host_smc_handler))
 		handled = default_host_smc_handler(host_ctxt);
 	if (!handled)
 		__kvm_hyp_host_forward_smc(host_ctxt);
diff --git a/arch/arm64/kvm/hyp/nvhe/psci-relay.c b/arch/arm64/kvm/hyp/nvhe/psci-relay.c
index d4825b6..f8db544 100644
--- a/arch/arm64/kvm/hyp/nvhe/psci-relay.c
+++ b/arch/arm64/kvm/hyp/nvhe/psci-relay.c
@@ -28,14 +28,19 @@
 static void (*pkvm_psci_notifier)(enum pkvm_psci_notification, struct kvm_cpu_context *);
 static void pkvm_psci_notify(enum pkvm_psci_notification notif, struct kvm_cpu_context *host_ctxt)
 {
-	if (READ_ONCE(pkvm_psci_notifier))
+	if (smp_load_acquire(&pkvm_psci_notifier))
 		pkvm_psci_notifier(notif, host_ctxt);
 }
 
 #ifdef CONFIG_MODULES
 int __pkvm_register_psci_notifier(void (*cb)(enum pkvm_psci_notification, struct kvm_cpu_context *))
 {
-	return cmpxchg(&pkvm_psci_notifier, NULL, cb) ? -EBUSY : 0;
+	/*
+	 * Paired with smp_load_acquire(&pkvm_psci_notifier) in
+	 * pkvm_psci_notify(). Ensure memory stores hapenning during a pKVM module
+	 * init are observed before executing the callback.
+	 */
+	return cmpxchg_release(&pkvm_psci_notifier, NULL, cb) ? -EBUSY : 0;
 }
 #endif
 
diff --git a/arch/arm64/kvm/hyp/nvhe/serial.c b/arch/arm64/kvm/hyp/nvhe/serial.c
index 0b2cf3b..475ebf4 100644
--- a/arch/arm64/kvm/hyp/nvhe/serial.c
+++ b/arch/arm64/kvm/hyp/nvhe/serial.c
@@ -35,7 +35,8 @@
 
 static inline bool hyp_serial_enabled(void)
 {
-	return !!READ_ONCE(__hyp_putc);
+	/* Paired with __pkvm_register_serial_driver()'s cmpxchg */
+	return !!smp_load_acquire(&__hyp_putc);
 }
 
 void hyp_puts(const char *s)
@@ -64,5 +65,10 @@
 
 int __pkvm_register_serial_driver(void (*cb)(char))
 {
-	return cmpxchg(&__hyp_putc, NULL, cb) ? -EBUSY : 0;
+	/*
+	 * Paired with smp_load_acquire(&__hyp_putc) in
+	 * hyp_serial_enabled(). Ensure memory stores hapenning during a pKVM
+	 * module init are observed before executing the callback.
+	 */
+	return cmpxchg_release(&__hyp_putc, NULL, cb) ? -EBUSY : 0;
 }